#include "server.h"
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <netinet/in.h>
#include <optional>
#include <pthread.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
head_buf_t *malloc_head_buf(int head_len)
{
    char *buf = (char *)malloc(sizeof(char) * head_len);
    if (!buf)
    {
        printf("malloc head buf failed\n");
        return NULL;
    }
    head_buf_t *head_buf = (head_buf_t *)malloc(sizeof(head_buf_t));
    if (!head_buf)
    {
        free(buf);
        printf("malloc head buf failed\n");
        return NULL;
    }
    head_buf->buf = buf;
    head_buf->data_len = head_len;
    head_buf->offset = 0;
    return head_buf;
}
void dealloc_head_buf(head_buf_t *head_buf)
{
    if (!head_buf)
    {
        return;
    }
    if (head_buf->buf)
    {
        free(head_buf->buf);
    }
    free(head_buf);
}
data_buf_t *malloc_data_buf(int type_id, int data_len)
{
    char *buf = (char *)malloc(sizeof(char) * data_len);
    if (!buf)
    {
        printf("malloc data buf failed\n");
        return NULL;
    }
    data_buf_t *data_buf = (data_buf_t *)malloc(sizeof(data_buf_t));
    if (!data_buf)
    {
        free(buf);
        printf("malloc data buf failed\n");
        return NULL;
    }
    data_buf->buf = buf;
    data_buf->type = type_id;
    data_buf->data_len = data_len;
    data_buf->offset = 0;
    return data_buf;
}
void dealloc_data_buf(data_buf_t *data_buf)
{
    if (!data_buf)
    {
        return;
    }

    if (data_buf->buf)
    {
        free(data_buf->buf);
    }

    free(data_buf);
}
send_node_t *malloc_send_node(int type_id, int data_len)
{
    send_node_t *send_node = (send_node_t *)malloc(sizeof(send_node_t));
    if (!send_node)
    {
        printf("malloc send_node failed\n");
        return NULL;
    }

    data_buf_t *send_buf = malloc_data_buf(type_id, data_len);
    if (!send_buf)
    {
        free(send_node);
        printf("malloc send_node failed\n");
        return NULL;
    }

    send_node->send_buf = send_buf;
    send_node->next = NULL;

    return send_node;
}
void dealloc_send_node(send_node_t *send_node)
{
    if (!send_node)
    {
        return;
    }
    if (send_node->send_buf)
    {
        dealloc_data_buf(send_node->send_buf);
    }

    free(send_node);
}
send_que_t *malloc_send_que()
{
    send_node_t *head = malloc_send_node(0, BUFF_SIZE);
    if (!head)
    {
        printf("malloc send que failed\n");
        return NULL;
    }
    send_que_t *send_que = (send_que_t *)malloc(sizeof(send_que_t));
    if (!send_que)
    {
        dealloc_send_node(head);
        printf("malloc send que failed!\n");
        return NULL;
    }
    send_que->head = head;
    if (0 != pthread_mutex_init(&send_que->send_que_lock, NULL))
    {
        fprintf(stderr, "mutex init failed!\n");
        dealloc_send_node(head);
        free(send_que);
        return NULL;
    }
    return send_que;
}
void dealloc_send_que(send_que_t *send_que)
{
    if (!send_que)
    {
        return;
    }
    if (pthread_mutex_lock(&send_que->send_que_lock) != 0)
    {
        fprintf(stderr, "failed to lock mutex in free_send_que\n");
        // 虽然锁失败，但我们仍尝试清理资源
    }
    send_node_t *cur = send_que->head;
    while (cur)
    {
        send_node_t *next = cur->next;
        dealloc_send_node(cur);
        cur = next;
    }
    send_que->head = NULL;
    if (pthread_mutex_unlock(&send_que->send_que_lock) != 0)
    {
        fprintf(stderr, "failed to unlock mutex in free_send_que\n");
    }
    // 4) 销毁互斥锁
    if (pthread_mutex_destroy(&send_que->send_que_lock) != 0)
    {
        fprintf(stderr, "failed to destroy mutex in free_send_que\n");
    }

    // 5) 释放队列结构体
    free(send_que);
}
session_t *session_create(int fd)
{
    session_t *session = (session_t *)malloc(sizeof(session_t));
    if (!session)
    {
        printf("failed to create session\n");
        return NULL;
    }
    session->fd = fd;
    session->data_buf = NULL;
    session->head_buf = malloc_head_buf(HEAD_LEN);
    session->send_que = malloc_send_que();
    session->recv_stage = NO_RECV;
    if (!(session->send_que))
    {
        printf("failed to create session send_que\n");
        free(session);
        return NULL;
    }

    return session;
}
void dealloc_session(session_t *session)
{
    if (session->head_buf)
    {
        dealloc_head_buf(session->head_buf);
    }

    if (session->data_buf)
    {
        dealloc_data_buf(session->data_buf);
    }

    if (session->send_que)
    {
        dealloc_send_que(session->send_que);
    }

    free(session);
}
void session_destroy(session_t *conn, int epoll_fd)
{
    epoll_ctl(epoll_fd, EPOLL_CTL_DEL, conn->fd, NULL);
    close(conn->fd);
    dealloc_session(conn);
}
int set_nonblocking(int fd)
{
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1)
        return -1;
    return fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

int create_and_bind()
{
    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd < 0)
    {
        perror("socket");
        exit(1);
    }
    int opt = 1;
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(SERVER_PORT);
    if (bind(listen_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0)
    {
        perror("bind");
        exit(1);
    }
    if (listen(listen_fd, SOMAXCONN) < 0)
    {
        perror("listen");
        exit(1);
    }
    return listen_fd;
}
void enqueue_node(send_que_t *q, send_node_t *node, int epoll_fd, int fd, uint32_t evs)
{
    send_node_t *p = q->head;
    while (p->next)
    {
        p = p->next;
    }
    p->next = node;
    if (evs & EPOLLOUT)
    {
        return;
    }
    struct epoll_event ev;
    ev.events = EPOLLIN | EPOLLET | EPOLLOUT;
    ev.data.fd = fd;
    epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &ev);
}
ssize_t async_send(session_t *sess, int epoll_fd, uint32_t evs, size_t msg_type, size_t body_len, char *buf)
{
    if (!sess || !sess->send_que)
    {
        return -1;
    }
    send_node_t *send_node = malloc_send_node(msg_type, body_len);
    if (!send_node)
    {
        printf("send_node allocate failed!\n");
        return -1;
    }
    uint16_t net_type = htons(msg_type);
    uint16_t net_len = htons(body_len);
    memcpy(send_node->send_buf->buf, &net_type, 2);
    memcpy(send_node->send_buf->buf + 2, &net_len, 2);
    memcpy(send_node->send_buf->buf + 4, buf, body_len);
    size_t remain = send_node->send_buf->offset;
    // 加锁，保证线程安全
    if (pthread_mutex_lock(&sess->send_que->send_que_lock) != 0)
    {
        dealloc_send_node(send_node);
        fprintf(stderr, "Failed to lock mutex in send_data\n");
        return -1;
    }
    // 队列为空就直接发送
    if (sess->send_que->head->next == NULL)
    {
        ssize_t total_sent = 0;
        while (send_node->send_buf->offset < send_node->send_buf->data_len)
        {
            ssize_t result = write(sess->fd,
                                   send_node->send_buf->buf + send_node->send_buf->offset,
                                   send_node->send_buf->data_len - send_node->send_buf->offset);
            if (result > 0)
            {
                send_node->send_buf->offset += result;
                total_sent += result;
                continue;
            }
            if (errno == EAGAIN || errno == EWOULDBLOCK)
            {
                break;
            }
            // error
            printf("perr close\n");
            dealloc_send_node(send_node);
            // 解锁并返回
            pthread_mutex_unlock(&sess->send_que->send_que_lock);
            return -1;
        }
        if (send_node->send_buf->offset == send_node->send_buf->data_len)
        {
            // 数据发送完全
            dealloc_send_node(send_node);
            // 解锁并返回
            pthread_mutex_unlock(&sess->send_que->send_que_lock);
            return total_sent;
        }
        // 没发完全，节点入队
        enqueue_node(sess->send_que, send_node, epoll_fd, sess->fd, evs);
        // 解锁并返回
        pthread_mutex_unlock(&sess->send_que->send_que_lock);
        return total_sent;
    }
    // 队列不为空说明有其他线程发送EAGAIN了，那就放入队列
    enqueue_node(sess->send_que, send_node, epoll_fd, sess->fd, evs);
    pthread_mutex_unlock(&sess->send_que->send_que_lock);
    return 0;
}
// Dequeue and send on EPOLLOUT
void handle_epollout(int epoll_fd, int fd, send_que_t *q)
{
    pthread_mutex_lock(&q->send_que_lock);
    struct epoll_event ev;
    send_node_t *prev = q->head;
    send_node_t *cur = prev->next;
    while (cur)
    {
        data_buf_t *b = cur->send_buf;
        ssize_t n;
        while (b->offset < b->data_len)
        {
            n = write(fd, b->buf + b->offset, b->data_len - b->offset);
            if (n > 0)
            {
                b->offset += n;
                continue;
            }
            if (errno == EAGAIN || errno == EWOULDBLOCK)
            {
                // kernel buffer full -> wait next EPOLLOUT
                ev.events = EPOLLIN | EPOLLET | EPOLLOUT;
                ev.data.fd = fd;
                epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &ev);
                pthread_mutex_unlock(&q->send_que_lock);
                return;
            }
            // unrecoverable error: close connection
            close(fd);
            pthread_mutex_unlock(&q->send_que_lock);
            return;
        }
        // fully sent: remove node
        prev->next = cur->next;
        dealloc_send_node(cur);
        cur = prev->next;
    }
    // queue empty: disable EPOLLOUT
    ev.events = EPOLLIN | EPOLLET;
    ev.data.fd = fd;
    epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &ev);
    pthread_mutex_unlock(&q->send_que_lock);
}
static int event_count = 1024;
// 全局标志，信号处理后置为1
static volatile sig_atomic_t stop_server = 0;
// 信号处理器：收到 SIGINT/SIGTERM 时设置退出标志
void handle_signal(int signo) { stop_server = 1; }
int main(int argc, char *argv[])
{
    // 1. 注册信号处理
    struct sigaction sa = {0};
    sa.sa_handler = handle_signal;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = 0;               // 无特殊标志
    sigaction(SIGINT, &sa, NULL);  // 注册 SIGINT 的处理
    sigaction(SIGTERM, &sa, NULL); // 注册 SIGTERM 的处理
    std::map<int, session_t *> session_map;

    int listen_fd = create_and_bind();
    set_nonblocking(listen_fd);
    int epoll_fd = epoll_create1(0);
    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = listen_fd;
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_fd, &ev);
    // 构造事件数组
    struct epoll_event *event_addr = (struct epoll_event *)malloc(sizeof(struct epoll_event) * event_count);
    while (!stop_server)
    {
        int nfds = epoll_wait(epoll_fd, event_addr, event_count, -1);
        if (nfds < 0)
        {
            if (errno == EINTR)
                continue; // 被信号打断，重试 epoll_wait
            perror("epoll_wait");
            break;
        }
        if ((size_t)nfds == event_count)
        {
            size_t new_count = event_count * 2;
            struct epoll_event *new_addr =
                (struct epoll_event *)realloc(event_addr, sizeof(struct epoll_event) * new_count);
            if (!new_addr)
            {
                // realloc 失败时继续用旧的数组，或根据需要退出
                perror("realloc event_addr");
            }
            else
            {
                event_addr = new_addr;
                event_count = new_count;
                printf(">>> expanded event_addr to %d entries\n", event_count);
            }
        }
        for (int i = 0; i < nfds; i++)
        {
            int fd = event_addr[i].data.fd;
            uint32_t evs = event_addr[i].events;
            if (fd == listen_fd)
            {
                // TODO 编写获取新链接的函数
                struct sockaddr_in cli;
                socklen_t len = sizeof(cli);
                while (1)
                {
                    int conn_fd = accept(listen_fd, (struct sockaddr *)&cli, &len);
                    if (conn_fd < 0)
                    {
                        if (errno == EAGAIN || errno == EWOULDBLOCK)
                            break;
                        if (errno == EINTR)
                            continue;
                        perror("accept");
                        break;
                    }
                    // 创建session
                    session_t *session = session_create(conn_fd);
                    if (!session)
                    {
                        perror("session create failed!\n");
                        break;
                    }
                    set_nonblocking(conn_fd);
                    ev.events = EPOLLIN | EPOLLET;
                    ev.data.fd = conn_fd;
                    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, conn_fd, &ev);
                    session_map[conn_fd] = session;
                    printf("Accepted fd=%d\n", conn_fd);
                }
                continue;
            }
            auto it = session_map.find(fd);
            session_t *sess;
            if (it != session_map.end())
            {
                sess = it->second;
            }
            else
            {
                printf("can't found fd %d\n", fd);
                epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, NULL);
                continue;
            }
            if (evs & (EPOLLERR | EPOLLHUP))
            {
                // todo  错误处理
                continue;
            }
            // 有读或者写事件
            if (evs & (EPOLLIN | EPOLLOUT))
            {
                if (evs & EPOLLIN)
                {
                    // todo 编写读取数据逻辑
                    while (1)
                    {
                        if (sess->recv_stage == NO_RECV || sess->recv_stage == HEAD_RECVING)
                        {
                            // 获取剩余需要读取的数据长度
                            int remain = HEAD_LEN - sess->head_buf->offset;
                            ssize_t read_len = read(fd, sess->head_buf->buf + sess->head_buf->offset, remain);
                            if (read_len < 0)
                            {
                                if (errno == EAGAIN || errno == EWOULDBLOCK)
                                {
                                    // 数据已经读完
                                    break;
                                }

                                if (errno == EINTR)
                                {
                                    // 被信号打断，重试 read
                                    continue;
                                }
                                // 真正的错误
                                perror("read header");
                                goto net_error;
                            }
                            // 对端关闭
                            if (read_len == 0)
                            {
                                goto net_error;
                            }
                            // 正常读取数据
                            if (read_len < remain)
                            {
                                // 未读出完整头部
                                sess->recv_stage = HEAD_RECVING;
                                sess->head_buf->offset += read_len;
                                continue;
                            }
                            // 读完整了头部
                            // 处理头部信息
                            uint8_t *hdr = (uint8_t *)sess->head_buf->buf;
                            uint16_t t, l;
                            memcpy(&t, hdr, sizeof(t));
                            memcpy(&l, hdr + 2, sizeof(l));
                            uint16_t msg_type = ntohs(t);
                            uint16_t body_len = ntohs(l);

                            // 检验对方发送程度字段是否合理
                            if (body_len > BUFF_SIZE)
                            {
                                printf("msg body too big, throw!\n");
                                session_destroy(sess, epoll_fd);
                                session_map.erase(fd);
                                break;
                            }

                            // 数据正常
                            printf("receiv head id is %d\n", msg_type);
                            printf("receiv head data len is %d", body_len);
                            sess->recv_stage = BODY_RECVING;
                            sess->head_buf->offset += read_len;
                            sess->data_buf = malloc_data_buf(msg_type, body_len);
                        }
                        while (sess->recv_stage == BODY_RECVING)
                        {
                            // 剩余要读取的长度
                            int remain = sess->data_buf->data_len - sess->data_buf->offset;
                            ssize_t read_len = read(fd, sess->data_buf->buf + sess->data_buf->offset, remain);
                            if (read_len < 0)
                            {
                                if (errno == EAGAIN || errno == EWOULDBLOCK)
                                {
                                    // 数据已读完
                                    break;
                                }
                                if (errno == EINTR)
                                {
                                    // 信号被打断，重试read
                                    continue;
                                }

                                perror("read header");
                                session_destroy(sess, epoll_fd);
                                session_map.erase(fd);
                                break;
                            }
                            // 对端关闭
                            if (read_len == 0)
                            {
                                session_destroy(sess, epoll_fd);
                                session_map.erase(fd);
                                break;
                            }
                            sess->data_buf->offset += read_len;
                            if (read_len < remain)
                            {
                                continue;
                            }
                            // TODO 处理消息
                            uint16_t body_len = sess->data_buf->data_len;
                            uint16_t msg_type = sess->data_buf->type;
                            size_t total_len = sizeof(uint16_t *) * 2 + body_len;
                            uint8_t *buf = (uint8_t *)malloc(total_len);
                            if (buf)
                            {
                                uint16_t net_type = htons(msg_type);
                                uint16_t net_len = htons(body_len);
                                memcpy(buf, &net_type, 2);
                                memcpy(buf + 2, &net_len, 2);
                                memcpy(buf + 4, sess->data_buf->buf, body_len);
                                write(fd, buf, total_len);
                                free(buf);
                            }
                            // 回收消息
                            dealloc_data_buf(sess->data_buf);
                            sess->data_buf = NULL;
                            sess->recv_stage = NO_RECV;
                            memset(sess->head_buf->buf, 0, HEAD_LEN);
                            sess->head_buf->offset = 0;
                            continue;
                        }
                    }
                }

                if (evs & EPOLLOUT)
                {
                    // todo 编写写数据逻辑
                }
                continue;
            }
        net_error:
            // 3) 其它
            fprintf(stderr, "fd=%d unexpected event: -1x%x\n", fd, evs);
            session_destroy(sess, epoll_fd);
            it = session_map.find(fd);
            if (it != session_map.end())
            {
                session_map.erase(it); // 无额外查找开销
            }
        }
    }
    // 遍历回收内存
    for (const auto pair: session_map)
    {
        auto sess = pair.second;
        if (sess)
        {
            session_destroy(sess, epoll_fd);
        }
    }
    session_map.clear();
    free(event_addr);
    return 0;
}
